from camel.agents import ChatAgent
from camel.messages import BaseMessage
from camel.models import ModelFactory
from camel.toolkits import MCPToolkit
from camel.types import ModelPlatformType, RoleType
from dotenv import load_dotenv
import asyncio
import os

from StructureGraph_Construction import connect_mysql, connect_neo4j, get_table_structure, MYSQL_CONFIG, NEO4J_CONFIG

load_dotenv()
Model_Type = os.getenv("MODEL_TYPE")
Model_Api = os.getenv("ZHIPU_API_KEY")
Model_Url = os.getenv("MODEL_URL")
model = ModelFactory.create(
    model_platform=ModelPlatformType.OPENAI,
    model_type=Model_Type,
    api_key=Model_Api,
    url=Model_Url
)

FAULT_TYPE_NAME=['空指针异常','数组越界异常','类型转换异常','并发修改异常','栈溢出错误','内存溢出错误','数字格式异常','非法参数异常','非法状态异常','类找不到异常']

def aim_extraction(question):
    """从文本中提取目标信息"""
    # 在这里实现你的目标信息提取逻辑
    sys_mes = f"""
        你是一个信息提取专家，专门从用户提供的查询中提取关键信息。请根据以下要求严格执行：

        【目标】
        从用户问题中提取出与故障相关的关键信息。考虑问题的上下文，综合分析问题的背景、错误描述和可能的原因。

        【提取信息要求】
        1. **故障类型**：例如“空指针异常”、“数据库连接失败”等。根据问题推断出最可能的故障类型,从{FAULT_TYPE_NAME}中选择最相关的一个，如果毫无关联性则默认不填。
        2. **故障描述**：例如“在第15行调用User对象的getUsername方法时出现错误”或“数据库查询时发生超时”等。详细描述故障发生的环境和表现。
        3. **故障原因**：例如“User对象为null”或“数据库连接池已满”等。根据上下文推测故障的可能根本原因。
        4. **错误代码**：如果问题中提到了错误代码，如“Exception in thread 'main' java.lang.NullPointerException: null”，请提取并标注出来。
        5. **错误位置**：如果问题中提到了具体的行号、函数或模块，请提取错误位置的信息（如：“第15行”、“myFunction()”等）。
        6. **时间信息**：如果问题涉及到发生时间或周期性问题，请提取相关时间信息（如：“今天”、“上午”、“每次启动时”）。

        你的输出应为一个字典，包含以下字段：
        - 'fault_type': 故障类型（例如："NullPointerException"）
        - 'fault_description': 故障描述（例如："在第15行调用User对象的getUsername方法时出现错误"）
        - 'root_cause': 故障原因（例如："User对象为null"）
        - 'error_code': 错误代码（如果有，可能为空或null）
        - 'error_location': 错误位置（如果有，可能为空或null）
        - 'time_info': 时间信息（如果有，可能为空或null）

        **注意事项**：
        - 若未能识别某一项信息，可以返回null或空字符串。
        - 请不要输出任何多余的信息，只需返回格式化后的字典结果。
        """

    Agent_Text = ChatAgent(
        model=model,
        output_language="zh",
        system_message=sys_mes,
    )
    mes = BaseMessage(
        role_name='User',
        role_type=RoleType.USER,
        content=question,
        meta_dict={}
    )
    response = Agent_Text.step(mes)
    print(f"\n\n用户问题提取智能体回复：\n{response.msg.content}")
    return response.msg.content


def query_graph_relationships(graph, table_names=None):
    """查询Neo4j中表之间的关系,
    这个函数的作用是从Neo4j图数据库中查询指定表之间的关系信息，为Text2SQL系统提供表间关联的语义信息。
    """
    if not table_names:  # 如果未指定要查询的表名，则使用默认表名列表
        table_names = ["fault_info", "fault_solution", "fault_type"]

    query = """
    MATCH (t1:Table)-[r:REFERENCES]-(t2:Table)
    WHERE t1.name IN $table_names AND t2.name IN $table_names AND t1.name < t2.name
    RETURN t1.name AS from_table, r.field_relation AS relationship, t2.name AS to_table
    """

    result = graph.run(query, table_names=table_names).data()
    print('--->成功获取到Neo4j中表之间的关系<---')
    return result


def query_fault_data(graph, query_type="all", **kwargs):
    """查询故障数据知识图谱
    
    Args:
        graph: Neo4j图数据库连接
        query_type: 查询类型 ("all", "by_type", "by_symptom", "by_id", "statistics")
        **kwargs: 查询参数
    
    Returns:
        查询结果列表
    """
    try:
        if query_type == "all":
            # 查询所有故障类型
            query = """
            MATCH (ft:FaultType)
            RETURN ft.type_id AS type_id, ft.name AS name, ft.description AS description
            ORDER BY ft.type_id
            """
            result = graph.run(query).data()

        elif query_type == "by_type":
            # 根据故障类型查询故障信息
            type_name = kwargs.get('type_name', '')
            query = """
            MATCH (ft:FaultType {name: $type_name})<-[:HAS_TYPE]-(fi:FaultInfo)
            RETURN fi.fault_id AS fault_id, fi.category AS category, 
                   fi.symptom AS symptom, fi.root_cause AS root_cause
            """
            result = graph.run(query, type_name=type_name).data()

        elif query_type == "by_symptom":
            # 根据症状关键词搜索故障
            symptom_keyword = kwargs.get('symptom_keyword', '')
            query = """
            MATCH (fi:FaultInfo)-[:HAS_TYPE]->(ft:FaultType)
            WHERE fi.symptom CONTAINS $symptom_keyword
            RETURN fi.fault_id AS fault_id, fi.category AS category, 
                   fi.symptom AS symptom, ft.name AS fault_type
            """
            result = graph.run(query, symptom_keyword=symptom_keyword).data()

        elif query_type == "by_id":
            # 根据故障ID获取完整信息
            fault_id = kwargs.get('fault_id', 0)
            query = """
            MATCH (ft:FaultType)<-[:HAS_TYPE]-(fi:FaultInfo {fault_id: $fault_id})-[:HAS_SOLUTION]->(fs:FaultSolution)
            RETURN ft.name AS fault_type, ft.description AS type_description,
                   fi.category AS category, fi.symptom AS symptom, fi.root_cause AS root_cause,
                   fs.solution_id AS solution_id, fs.solution AS solution
            """
            result = graph.run(query, fault_id=fault_id).data()

        elif query_type == "statistics":
            # 获取故障统计信息
            query = """
            MATCH (ft:FaultType)<-[:HAS_TYPE]-(fi:FaultInfo)
            RETURN ft.name AS fault_type, COUNT(fi) AS fault_count
            ORDER BY fault_count DESC
            """
            result = graph.run(query).data()

        else:
            result = []

        print(f'--->成功查询故障数据，类型: {query_type}，结果数量: {len(result)}<---')
        return result

    except Exception as e:
        print(f'--->故障数据查询失败: {e}<---')
        return []


def Agent_Graph(keyWords, question):
    def get_table_structure_tool():
        """
        工具描述：获取所有表的结构信息，包括表名、字段名、字段类型、主键信息等。
        返回值：以字典形式返回所有表的结构详情，适用于数据库结构理解与SQL生成。
        """
        return get_table_structure(conn)

    def query_graph_relationships_tool(table_names=None):
        """
        工具描述：查询指定表之间的关系（如外键、业务关联等），如未指定表名则返回所有表关系。
        参数：table_names（可选）- 需要查询关系的表名列表。
        返回值：返回表与表之间的关系详情，便于理解数据表之间的业务联系。
        """
        return query_graph_relationships(graph, table_names)

    def query_all_fault_types_tool():
        """
        工具描述：获取所有故障类型信息，包括类型ID、名称和描述。
        返回值：返回所有故障类型的详细信息。
        """
        return query_fault_data(graph, "all")

    def query_faults_by_type_tool(type_name):
        """
        工具描述：根据故障类型名称查询相关故障信息。
        参数：type_name - 故障类型名称。
        返回值：返回指定类型的所有故障信息。
        """
        return query_fault_data(graph, "by_type", type_name=type_name)

    def query_faults_by_symptom_tool(symptom_keyword):
        """
        工具描述：根据症状关键词搜索相关故障信息。
        参数：symptom_keyword - 症状关键词。
        返回值：返回包含指定症状关键词的故障信息。
        """
        return query_fault_data(graph, "by_symptom", symptom_keyword=symptom_keyword)

    def query_fault_by_id_tool(fault_id):
        """
        工具描述：根据故障ID获取完整的故障信息，包括类型、症状、原因和解决方案。
        参数：fault_id - 故障ID。
        返回值：返回指定故障的完整详细信息。
        """
        return query_fault_data(graph, "by_id", fault_id=fault_id)

    def query_fault_statistics_tool():
        """
        工具描述：获取故障统计信息，包括各类型故障的数量分布。
        返回值：返回故障类型统计信息。
        """
        return query_fault_data(graph, "statistics")

    graph = connect_neo4j()
    conn = connect_mysql()
    agent_graph = ChatAgent(
        model=model,
        output_language="zh",
        system_message=f"""
            你是一个智能知识图谱查询助手，必须严格按照以下执行流程操作：

            【工具功能】
            1. **get_table_structure_tool**: 获取所有表的结构信息，包括表名、字段名、字段类型、主键信息等。此工具帮助理解数据库结构，生成SQL查询时需要用到。
            2. **query_graph_relationships_tool**: 查询指定表之间的关系，例如外键关联、业务关联等。如果用户的查询涉及到多个表之间的关系，你需要使用此工具获取表与表之间的关联信息。
            3. **query_all_fault_types_tool**: 获取所有故障类型信息，包括类型ID、名称和描述。此工具适用于查询所有故障类型信息的需求。
            4. **query_faults_by_type_tool**: 根据故障类型名称查询相关故障信息。此工具适用于根据用户提供的故障类型进行数据查询。
            5. **query_faults_by_symptom_tool**: 根据症状关键词查询相关故障信息。如果用户查询的是某种故障症状，使用此工具进行查询。
            6. **query_fault_by_id_tool**: 根据故障ID获取完整的故障信息，包括类型、症状、原因和解决方案。如果用户提供故障ID，使用此工具来获取详细信息。
            7. **query_fault_statistics_tool**: 获取故障统计信息，例如各类型故障的数量分布。如果用户希望查看故障类型的统计数据，使用此工具。

            【执行流程】
            1. **获取表结构**：首先调用 `get_table_structure_tool()` 获取完整的表结构信息。理解数据库表结构对于后续的 SQL 查询和知识图谱查询至关重要。
            2. **获取表关系**：然后调用 `query_graph_relationships_tool()` 获取表与表之间的关系。如果用户查询需要跨多个表的数据，应确保知道这些表之间的关系。
            3. **根据问题选择工具**：根据用户的问题，从上述工具中选择一个或多个进行调用。你必须根据问题类型来判断需要使用哪个工具进行查询：
                - 如果问题涉及到**故障类型**，使用 `query_all_fault_types_tool` 或 `query_faults_by_type_tool`。
                - 如果问题涉及到**故障症状**，使用 `query_faults_by_symptom_tool`。
                - 如果问题提供了**故障ID**，使用 `query_fault_by_id_tool`。
                - 如果问题要求查看**故障统计**，使用 `query_fault_statistics_tool`。
            4. **返回查询结果**：将查询结果完整返回，并且不能修改字段名或表名，确保输出的数据与工具返回的数据一致。

            【严格要求】
            - 必须使用工具返回的完整字段名，不得简写或省略。
            - 不得自行生成或修改字段名。
            - 必须保持工具返回的原始数据结构。
            - 输出格式：表结构信息 + 表关系信息 + 数据查询结果。

            请根据以下用户提供的问题，并将关键字作为参考，调用合适的工具进行查询并返回结果：
            关键字为：{keyWords}
            用户问题：{question}
            """,
        tools=[
            query_graph_relationships_tool,
            get_table_structure_tool,
            query_all_fault_types_tool,
            query_faults_by_type_tool,
            query_faults_by_symptom_tool,
            query_fault_by_id_tool,
            query_fault_statistics_tool
        ]
    )
    msg = BaseMessage(
        role_name="用户",
        role_type=RoleType.USER,
        content=f"""
    请你根据我提供的关键字进行查询，严格按照系统提示词进行输出。
    关键词为：{question}
""",
        meta_dict={}
    )
    response = agent_graph.step(msg)
    print(f"\n\n知识图谱查询智能体回复：\n{response.msg.content}")
    return f"知识图谱查询智能体回复：{response.msg.content}"


async def run_mcp_client(keyWords, question, graph_context=""):
    # 1. 初始化 MCPToolkit，指定配置文件路径
    mcp_toolkit = MCPToolkit(config_path="config/mcp_tools.json", timeout=1200)

    # 2. 异步连接到配置文件中定义的所有 MCP 服务器
    await mcp_toolkit.connect()

    # 3. 创建 CAMEL Agent 实例，并将从 MCPToolkit 获取的工具列表传递给 Agent
    camel_agent = ChatAgent(
        model=model,
        tools=[*mcp_toolkit.get_tools()],  # 动态获取并解包 MCP 工具列表
        system_message=f"""
            你是一个数据库专家，能够执行SQL语句，必须严格按照以下流程进行 SQL 查询：

            【执行流程】
            1. 首先获取并分析知识图谱上下文中的表结构信息，确保理解每个字段的详细含义。
            2. 结合 GraphAgent 返回的数据，确认字段信息和表结构。
            3. 如果知识图谱包含确切故障id，根据故障id查询其关联的信息，如果没有故障id，根据故障名称查询【若知识图谱上下文已经包含确切的主键信息，根据主键进行查询出准确信息即可，不需要进行模糊查询】
            4. 如果用户问题模糊以及知识图谱未查询到确切结果，则应提取用户问题，使用模糊查询（or）
            5. 根据用户提供的查询信息，结合知识图谱上下文，使用完整的表字段名生成一个 SQL 查询语句,应包含故障类型名称，故障名称，症状，导致原因，解决方法这五个字段（如果知识图谱已经包含准确的数据，按照知识图谱所返回的结果，如id等进行查询验证）。
            6. 将使用的SQL语句，以及查询的结果信息进行整合，以便于之后的智能体理解、使用。
            【重点】严格按照从知识图谱获取的表结构和字段信息，进行sql语句的查询，不能自己臆想字段、名称。
            
            【输出格式】
            1. SQL 查询语句：输出使用的 SQL 查询语句，必须是可执行的。
            2. 执行结果说明：简要描述执行 SQL 查询后的结果，提供查询结果的总结。
            3. 总结内容应包含指导用户解决问题的切实可行的执行方案（结合查询结果）
            
            
            
            用户问题：{question}
            提取的关键词：{keyWords}
            数据库配置：{MYSQL_CONFIG}
            知识图谱上下文：{graph_context}
            """
    )

    # 4. Agent 执行任务，可能调用 MCP 工具
    msg = BaseMessage(
        role_name="User",
        role_type=RoleType.USER,
        content=f"""
        请你严格按照系统提示词要求执行！
""",
        meta_dict={}
    )
    response = await camel_agent.astep(msg)
    await mcp_toolkit.disconnect()
    print(f"SQL查询智能体回复: \n{response.msg.content}")
    return f"SQL查询智能体回复: \n{response.msg.content}"


def Agent_SQL(keyWords, question):
    graph_result = Agent_Graph(keyWords, question)
    sql_result = asyncio.run(run_mcp_client(keyWords, question, graph_result))
    agent_sql = ChatAgent(
        model=model,
        system_message=f"""
        你是一个text2sql专家，你的任务是根据知识图谱上下文和SQL查询上下文生成 SQL 查询，并且输出规范的查询结果说明。，必须严格按照以下执行流程操作：
        
        【执行流程】
        1. 分析知识图谱上下文中的表结构、表间关系、字段、数据关系等信息，确保清晰理解数据库结构，以及数据间的关系。
        2. 分析SQL查询上下文中的SQL语句，以及查询后的结果。
        3. 严格按照工具返回的字段名生成SQL
        4. 验证SQL语句的可执行性
        
        【严格规定】
        - 必须使用知识图谱和数据查询中返回的完整字段名
        - 不得简写、省略或修改任何字段名
        - 不得自己生成字段名
        - 必须保持字段名的原始大小写和格式
        - 生成的SQL必须可执行
        
        【输出格式】[必须严格按照以下格式输出]

        1. 如果知识图谱上下文和SQL查询上下文中的信息不足以生成SQL，请输出“无法生成SQL”
        2. 生成的SQL语句（必须可执行），并且必须包含故障类型名称，故障名称，症状，导致原因，解决方法这五个字段。
        3. 执行结果说明
        4. 根据知识图谱上下文和数据查询信息，输出一份综合性结果
        5. 根据查询结果和综合性结果输出一份指导用户解决问题的切实可行的执行方案，应结构清晰，逻辑通顺。
        
        知识图谱上下文：{graph_result}
        SQL查询上下文：{sql_result}
        
        """,
        output_language="zh",
    )
    msg = BaseMessage(
        role_name='user',
        role_type=RoleType.USER,
        content=f"""
        请你严格按照系统提示词回答问题：{question}
""",
        meta_dict={}
    )
    response = agent_sql.step(msg)
    return f"\n汇总智能体回复: \n{response.msg.content}"


if __name__ == '__main__':
    question = """
    数组大小超过限制了，报错信息如下：
    Exception in thread "main" java.lang.OutOfMemoryError: Requested array size exceeds VM limit at LargeArrayDemo.createArray(LargeArrayDemo.java:15)
    怎么解决
    """
    keyWords = aim_extraction(question)
    print(Agent_SQL(keyWords, question))
